
#include "hash.h"
#define DRIVERTAG1 'HSAH'


static unsigned int NextPrime(int N)
{
	int i = 0;
	if (N % 2 == 0)
	{
		N++;
	}

	for (;; N += 2)
	{
		for (i = 3; i * i <= N; i += 2)
		{
			if (N % i == 0)
			{
				goto ContOuter;
			}
			return N;
		}
	ContOuter:;
	}
}

unsigned int Hash(DWORD key, unsigned int tableSize)
{
	return key % tableSize;
}

PHASHTABLE InitializeTable(unsigned int tableSize)
{
	PHASHTABLE		pHashTable = NULL;
	unsigned int	i = 0;
	NTSTATUS		status;

	pHashTable = (PHASHTABLE)ExAllocatePoolWithTag(PagedPool, sizeof(HASHTABLE), DRIVERTAG1);
	if (NULL == pHashTable)
	{
		KdPrint(("ExAllocatePoolWithTag returned NULL!\n"));
		goto err_ret;
	}

	pHashTable->tableSize = NextPrime(tableSize);

	// Allocate array of pointers to linkedlists 
	pHashTable->pListHeads = (PLIST_ENTRY)ExAllocatePoolWithTag(PagedPool, sizeof(LIST_ENTRY)* pHashTable->tableSize, DRIVERTAG1);
	if (pHashTable->pListHeads == NULL)
	{
		KdPrint(("ExAllocatePoolWithTag returned NULL!\n"));
		goto err_ret;
	}

	//// Allocate space for the lookaside list for the TWOWAY-structures.
	//pHashTable->LookasideList_TWOWAY = (PPAGED_LOOKASIDE_LIST)ExAllocatePoolWithTag(PagedPool, sizeof(NPAGED_LOOKASIDE_LIST), DRIVERTAG1);
	//if (pHashTable->pLookasideList_TWOWAY == NULL)
	//{
	//	DbgPrint("ExAllocatePoolWithTag returned NULL!\n");
	//	goto err_ret;
	//}

	// Initialize the lookaside list.
	ExInitializePagedLookasideList(
		&pHashTable->LookasideList_TWOWAY,
		NULL,
		NULL,
		0,
		sizeof(TWOWAY),
		DRIVERTAG1,
		0);

	// Allocate empty nodes for the each linked list.
	for (i = 0; i < pHashTable->tableSize; i++)
	{
		InitializeListHead(&pHashTable->pListHeads[i]);
	}
	status = ExInitializeResourceLite(&pHashTable->resource);
	if (!NT_SUCCESS(status))
	{
		goto err_ret;
	}
	return pHashTable;
err_ret:
	if (pHashTable != NULL)
	{
		ExDeletePagedLookasideList(&pHashTable->LookasideList_TWOWAY);
	}
	if (pHashTable->pListHeads != NULL)
	{
		ExFreePoolWithTag(pHashTable->pListHeads, DRIVERTAG1);
		pHashTable->pListHeads = NULL;
	}
	if (pHashTable != NULL)
	{
		ExFreePoolWithTag(pHashTable, DRIVERTAG1);
		pHashTable = NULL;
	}
	return NULL;
}

PTWOWAY Find(DWORD key, PHASHTABLE pHashTable)
{
	PTWOWAY pNode = NULL;
	PLIST_ENTRY pListHead = NULL;
	PLIST_ENTRY pListLink = NULL;

	LockRead(pHashTable);
	pListHead = pListLink = &pHashTable->pListHeads[Hash(key, pHashTable->tableSize)];
	if (pListHead == NULL)
	{
		KdPrint(("pListHead is NULL!\n"));
		UnLock(pHashTable);
		return NULL;
	}

	if (!IsListEmpty(pListHead))
	{
		pListLink = pListHead->Flink;
		do
		{
			pNode = CONTAINING_RECORD(pListLink, TWOWAY, linkfield);
			if (pNode->key == key)
			{
				UnLock(pHashTable);
				return pNode;
			}
			pListLink = pListLink->Flink;
		} while (pListLink != pListHead);
	}
	UnLock(pHashTable);
	return NULL;
}

BOOLEAN Insert(DWORD key, PELEMENT pData, PHASHTABLE pHashTable)
{
	PTWOWAY pNode = NULL;
	PTWOWAY pNewNode = NULL;
	PLIST_ENTRY pListHead = NULL;

	pNode = Find(key, pHashTable);
	if (pNode == NULL)
	{
		pNewNode = (PTWOWAY)ExAllocateFromPagedLookasideList(&pHashTable->LookasideList_TWOWAY);
		if (pNewNode == NULL)
		{
			KdPrint(("ExAllocateFromNPagedLookasideList returned NULL!\n"));
			return FALSE;
		}

		pNewNode->key = key;

		pNewNode->data.ulSize = pData->ulSize;
		pNewNode->data.pData = ExAllocatePoolWithTag(PagedPool, pData->ulSize, DRIVERTAG1);
		if (NULL == pNewNode->data.pData)
		{
			return FALSE;
		}
		RtlZeroMemory(pNewNode->data.pData, pData->ulSize);
		RtlMoveMemory(pNewNode->data.pData, pData->pData, pData->ulSize);
		LockWrite(pHashTable);
		pListHead = &pHashTable->pListHeads[Hash(key, pHashTable->tableSize)];
		InsertTailList(pListHead, &pNewNode->linkfield);
		UnLock(pHashTable);
		return TRUE;
	}
	else
	{
		Remove(key, pHashTable);
		Insert(key, pData, pHashTable);
		return TRUE;
	}
}

BOOLEAN Remove(DWORD key, PHASHTABLE pHashTable)
{
	PTWOWAY pNode = NULL;
	pNode = Find(key, pHashTable);

	if (pNode != NULL)
	{
		if (pNode->data.pData)
		{
			ExFreePool(pNode->data.pData);
		}
		LockWrite(pHashTable);
		RemoveEntryList(&pNode->linkfield);
		ExFreeToPagedLookasideList(&pHashTable->LookasideList_TWOWAY, pNode);
		UnLock(pHashTable);
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

void DestroyTable(PHASHTABLE pHashTable)
{
	PTWOWAY			pNode = NULL;
	PLIST_ENTRY		pListHead = NULL;
	PLIST_ENTRY		pListLink = NULL;
	unsigned int	i;
	NTSTATUS		status = STATUS_SUCCESS;

	LockWrite(pHashTable);
	if (pHashTable != NULL && pHashTable->pListHeads != NULL)
	{
		for (i = 0; i < pHashTable->tableSize; i++)
		{
			pListHead = pListLink = &pHashTable->pListHeads[i];
			if (pListHead == NULL)
			{
				KdPrint(("pListHead is NULL!\n"));
				continue;
			}
			if (!IsListEmpty(pListHead))
			{
				pListLink = pListHead->Flink;
				do
				{
					pNode = CONTAINING_RECORD(pListLink, TWOWAY, linkfield);
					if (pNode->data.pData)
					{
						ExFreePool(pNode->data.pData);
					}
					pListLink = pListLink->Flink;
					ExFreeToPagedLookasideList(&pHashTable->LookasideList_TWOWAY, pNode);
				} while (pListLink != pListHead);
			}
		}
	}

	if (pHashTable != NULL)
	{
		ExDeletePagedLookasideList(&pHashTable->LookasideList_TWOWAY);
		status = ExDeleteResourceLite(&pHashTable->resource);
		if (!NT_SUCCESS(status))
		{
			KdPrint(("ExDeleteResourceLite Failed!!\n"));
		}
	}
	if (pHashTable != NULL && pHashTable->pListHeads)
	{
		ExFreePoolWithTag(pHashTable->pListHeads, DRIVERTAG1);
	}

	UnLock(pHashTable);
	if (pHashTable)
	{
		ExFreePoolWithTag(pHashTable, DRIVERTAG1);
	}

}

ULONG  GetTableCount(PHASHTABLE pHashTable)
{
	PTWOWAY pNode = NULL;
	PLIST_ENTRY pListHead = NULL;
	PLIST_ENTRY pListLink = NULL;
	unsigned int i;
	ULONG total = 0;
	if (NULL == pHashTable)
	{
		return 0;
	}

	LockRead(pHashTable);
	if (pHashTable != NULL && pHashTable->pListHeads != NULL)
	{
		for (i = 0; i < pHashTable->tableSize; i++)
		{
			pListHead = pListLink = &pHashTable->pListHeads[i];
			if (pListHead == NULL)
			{
				KdPrint(("pListHead is NULL!\n"));
				continue;
			}
			if (!IsListEmpty(pListHead))
			{
				pListLink = pListHead->Flink;
				do
				{
					pNode = CONTAINING_RECORD(pListLink, TWOWAY, linkfield);
					pListLink = pListLink->Flink;
					if (pNode->key != 0)
					{
						total++;
					}
				} while (pListLink != pListHead);
			}
		}
	}
	UnLock(pHashTable);
	return total;
}

void LockWrite(PHASHTABLE pHashTable)
{
	KeEnterCriticalRegion();
	ExAcquireResourceExclusiveLite(&pHashTable->resource, TRUE);
}

void LockRead(PHASHTABLE pHashTable)
{
	KeEnterCriticalRegion();
	ExAcquireResourceSharedLite(&pHashTable->resource, TRUE);
}

void UnLock(PHASHTABLE pHashTable)
{
	ExReleaseResourceLite(&pHashTable->resource);
	KeLeaveCriticalRegion();
}